#!/usr/bin/python3

# Stellarium Web Engine - Copyright (c) 2022 - Stellarium Labs SRL
#
# This program is licensed under the terms of the GNU AGPL v3, or
# alternatively under a commercial licence.
#
# The terms of the AGPL v3 license can be found in the main directory of this
# repository.

# Script used to generate the file src/hip.inl, that contains a mapping of
# HIP -> healpix for fast HIP stars lookup.

# Note: originally we only used hipparcos new reduction, but some stars
# were missing so now we use both the original and the new reduction
# data!

from math import *

import csv
import functools
import hashlib
import healpy
import os
import sys


def hms_to_deg(hms):
    if hms.startswith("b'"): hms = hms[2:-2] # For gaia python3 bug.
    hms = [float(x) for x in hms.split()]
    return hms[0] * 15 + hms[1] * 15 / 60 + hms[2] * 15 / 60 / 60

def dms_to_deg(dms):
    if dms.startswith("b'"): dms = dms[2:-2] # For gaia python3 bug.
    dms = [float(x) for x in dms.split()]
    return dms[0] + dms[1] / 60 + dms[2] / 60 / 60

def generator(target, md5):
    """Decorator that checks if a generated file is already up to date
       This can be used to cache data that is slow to generate.
    """
    def decorator(func):
        @functools.wraps(func)
        def wrapper():
            path = os.path.join('/tmp/', target)
            if not os.path.exists(path):
                print('Generating %s' % path)
                func(path)
            current_md5 = hashlib.md5(open(path, 'rb').read()).hexdigest()
            if current_md5 != md5:
                print('Md5 for file %s changed!' % target)
                print('Current md5: %s' % current_md5)
                print('Expected   : %s' % md5)
                raise ValueError
            return path
        return wrapper
    return decorator


@generator('hip.csv', md5='f2f11a8f3f8ccae13f6eabf331b4e4db')
def get_hip_data_file(path):
    '''Get HIP data from gaia TAP server'''
    from astroquery.gaia import Gaia
    QUERY = '''
        SELECT hip, ra, de, vmag, rahms, dedms
        FROM public.hipparcos
        ORDER BY hip
    '''
    job = Gaia.launch_job_async(QUERY)
    r = job.get_results()
    r.write(path, overwrite=True, format='ascii.csv')
    if not hasattr(job, 'get_jobid'):
        print('Warning: cannot delete Gaia job')
        return
    Gaia.remove_jobs([job.get_jobid()])

@generator('hip_newreduction.csv', md5='b0b86439386bbcbae15efc402ed06da7')
def get_hip_newreduction_data_file(path):
    '''Get HIP data from gaia TAP server'''
    from astroquery.gaia import Gaia
    QUERY = '''
        SELECT hip, ra, dec, hp_mag AS vmag
        FROM public.hipparcos_newreduction
        ORDER BY hip
    '''
    job = Gaia.launch_job_async(QUERY)
    r = job.get_results()
    r.write(path, overwrite=True, format='ascii.csv')
    if not hasattr(job, 'get_jobid'):
        print('Warning: cannot delete Gaia job')
        return
    Gaia.remove_jobs([job.get_jobid()])

def run():
    htable = {}

    data = list(csv.DictReader(open(get_hip_newreduction_data_file())))
    for d in data:
        hip = int(d['hip'])
        ra = float(d['ra'])
        de = float(d['dec'])
        pix = healpy.ang2pix(4, ra, de, nest=True, lonlat=True)
        htable[hip] = pix

    # Add the stars missing from Hipparcos New Reduction!
    # XXX: would be nice to be able to load them all in a single SQL query.
    data = list(csv.DictReader(open(get_hip_data_file())))
    for d in data:
        hip = int(d['hip'])
        if hip in htable: continue
        ra = hms_to_deg(d['rahms'])
        de = dms_to_deg(d['dedms'])
        pix = healpy.ang2pix(4, ra, de, nest=True, lonlat=True)
        htable[hip] = pix

    data = [htable.get(x, 255) for x in range(max(htable.keys()))]

    out = open('src/hip.inl', 'w')
    print('// HIP -> pix pos at order 2.', file=out)
    print('// Generated by tools/make-hip-lookup.py', file=out)

    for i, v in enumerate(data):
        if (i % 15): print(' ', file=out, end='')
        print('%3d,' % v, file=out, end='')
        if (i + 1) % 15 == 0: print(file=out)

if __name__ == '__main__':
    run()
